package com.jsean.web;

import lombok.extern.slf4j.Slf4j;
import org.java_websocket.WebSocket;
import org.java_websocket.handshake.ClientHandshake;
import org.java_websocket.server.WebSocketServer;

import java.io.*;
import java.net.InetSocketAddress;
import java.util.Base64;
import java.util.Properties;

@Slf4j
public class FileTransferServer extends WebSocketServer {

    public FileTransferServer(int port) {
        super(new InetSocketAddress(port));
    }

    @Override
    public void onOpen(WebSocket conn, ClientHandshake handshake) {
        log.info("Client connected: " + conn.getRemoteSocketAddress());
    }

    @Override
    public void onClose(WebSocket conn, int code, String reason, boolean remote) {
        log.info("Client disconnected");
    }

    private String currentFileName;
    private long expectedFileSize;
    private long receivedBytes = 0;
    private FileOutputStream currentFileStream;
    private String KEY_WORD = "^^";
    private String SPLIT_WORD = "\\^\\^";
    private static final int CHUNK_SIZE = 1024 * 64; // 64KB chunks

    @Override
    public void onMessage(WebSocket conn, String message) {
        try {
            if (message.startsWith("FILE_METADATA" + KEY_WORD)) {
                // Handle file metadata
                String[] parts = message.split(SPLIT_WORD);
                currentFileName = new File(parts[1]).getName();
                expectedFileSize = Long.parseLong(parts[2]);
                receivedBytes = 0;
                
                // Create new file output stream
                File file = new File(currentFileName);
                currentFileStream = new FileOutputStream(file);
                
                conn.send("READY_FOR_CHUNKS:" + currentFileName);
                
            } else if (message.startsWith("CHUNK" + KEY_WORD)) {
                // Handle file chunks
                String[] parts = message.split(SPLIT_WORD);
                int chunkNumber = Integer.parseInt(parts[1]);
                byte[] chunkData = Base64.getDecoder().decode(parts[2]);
                
                currentFileStream.write(chunkData);
                receivedBytes += chunkData.length;
                
                // Send progress update
                if (chunkNumber % 10 == 0) {
                    double progress = (double)receivedBytes / expectedFileSize * 100;
                    conn.send("PROGRESS:" + currentFileName + ":" + String.format("%.2f", progress));
                }
                
            } else if (message.startsWith("FILE_END" + KEY_WORD)) {
                // Finalize file transfer
                currentFileStream.close();
                log.info("File received and saved as: " + currentFileName);
                conn.send("FILE_RECEIVED:" + currentFileName);
                
                // Reset state
                currentFileName = null;
                expectedFileSize = 0;
                receivedBytes = 0;
                currentFileStream = null;
                
            } else if (message.startsWith("DOWNLOAD_REQUEST" + KEY_WORD)) {
                log.info("download file");
                // Handle file download request
                String[] parts = message.split(SPLIT_WORD);
                String requestedFile = parts[1];
                File file = new File(requestedFile);
                
                if (file.exists()) {
                    // Send file metadata
                    String metadata = "FILE_METADATA" + KEY_WORD + requestedFile + KEY_WORD + file.length();
                    conn.send(metadata);
                    
                    // Send file chunks
                    try (FileInputStream fis = new FileInputStream(file)) {
                        byte[] buffer = new byte[CHUNK_SIZE];
                        int bytesRead;
                        int chunkNum = 0;
                        
                        while ((bytesRead = fis.read(buffer)) != -1) {
                            String chunk = "CHUNK" + KEY_WORD + chunkNum + KEY_WORD +
                                          Base64.getEncoder().encodeToString(buffer);
                            conn.send(chunk);
                            chunkNum++;
                        }
                        
                        // Send end marker
                        conn.send("FILE_END" + KEY_WORD + requestedFile);
                    }
                } else {
                    conn.send("ERROR: File not found: " + requestedFile);
                }
            } else {
                // Handle regular messages
                log.info("Received from client: " + message);
            }

        } catch (IOException e) {
            e.printStackTrace();
            try {
                if (currentFileStream != null) {
                    currentFileStream.close();
                }
            } catch (IOException ex) {
                ex.printStackTrace();
            }
            conn.send("ERROR:" + e.getMessage());
        }
    }

    @Override
    public void onError(WebSocket conn, Exception ex) {
        ex.printStackTrace();
    }

    @Override
    public void onStart() {
        log.info("Server started");
    }

    private static String getAppConfig() {
        Properties props = new Properties();
        try (InputStream is = FileTransferServer.class.getClassLoader().getResourceAsStream("application.properties")) {
            props.load(is);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        return props.getProperty("app.port", "8090");
    }

    public static void main(String[] args) {
        int port = Integer.parseInt(getAppConfig());
        FileTransferServer server = new FileTransferServer(port);
        server.start();
        log.info("Server is running on ws://localhost:8090");
    }
}